from common.HashMap import HashMap


def get_family_type_0(lookup, type):
    for item in lookup:
        if item in type:
            return True
    return False


def is_defined_type(type):

    if type is "Worms":
        return True
    if type is "Adware":
        return True
    if type is "Virus":
        return True
    if type is "Riskware":
        return True
    if type is "Spyware":
        return True
    if type is "Keylogger":
        return True
    if type is "Ransomware":
        return True
    if type is "Spam":
        return True
    if type is "Backdoor":
        return True
    if type is "Dropper":
        return True
    if type is "Downloader":
        return True
    if type is "Crypt":
        return True
    if type is "Agent":
        return True
    if type is "Rootkit":
        return True
    if type is "Trojan":
        return True
    if type is "Undefined":
        return True

    return False


def get_family_type(type):

    if type == "_":
        return None

    tmp_str = str(type).lower().replace(" ", "")

    if get_family_type_0(['worm', '[wrm]'], tmp_str) is True:
        return "Worms"

    if get_family_type_0(['adware'], tmp_str) is True:
        return "Adware"

    if get_family_type_0(['virus', 'expiro'], tmp_str) is True:
        return "Virus"

    if get_family_type_0(['riskware'], tmp_str) is True:
        return "Riskware"

    if get_family_type_0(['spyware', 'spy'], tmp_str) is True:
        return "Spyware"

    if get_family_type_0(['keylogger'], tmp_str) is True:
        return "Keylogger"

    if get_family_type_0(['ransom'], tmp_str) is True:
        return "Ransomware"

    if get_family_type_0(['spam'], tmp_str) is True:
        return "Spam"

    if get_family_type_0(['backdoor'], tmp_str) is True:
        return "Backdoor"

    if get_family_type_0(['dropper'], tmp_str) is True:
        return "Dropper"

    if get_family_type_0(['downloader'], tmp_str) is True:
        return "Downloader"

    if get_family_type_0(['crypt'], tmp_str) is True:
        return "Crypt"

    if get_family_type_0(['agent'], tmp_str) is True:
        return "Agent"

    if get_family_type_0(["trjoan",'trj/','trojan', 'troj', '[trj]'], tmp_str) is True:
        return "Trojan"

    if get_family_type_0(["rootkit"], tmp_str) is True:
        return "Rootkit"

    if get_family_type_0(['unsafe', 'unclassified', "unknown"], tmp_str) is True:
        return None

    return type


def choose_f_type(selected_choice):
    choose_type = ""
    count = 0
    for item in selected_choice.map:
        if item is not None:
            for tmp in item:
                if tmp[1] > count:
                    count = tmp[1]
                    choose_type = tmp[0]
    if not is_defined_type(choose_type):
        print(choose_type + " " + str(count))
        choose_type = "Undefined"
    return choose_type

def parse_and_determine_malware_type(line):
    selected_choice = HashMap()
    for cell in line.split(","):
        if cell is "":
            continue

        if "_" not in cell:
            continue

        index = cell.index("_")
        type = cell[index + 1:]

        family_type = get_family_type(type)
        if family_type is None:
            continue
        typeCount = selected_choice.get(family_type)
        if typeCount is not None:
            typeCount += 1
        else:
            typeCount = 1

        if family_type is not None:
            selected_choice.add(family_type, typeCount)

    return choose_f_type(selected_choice)


def parse_vt_data(file_name, type_map):
    file = open(file_name, "r")
    for line in file:
        line = line.strip()
        if line == "Error":
            continue

        index = line.index(":")
        hash_code = line[:index]
        line = line[index + 1:]

        choose_type = parse_and_determine_malware_type(line)
        # print(choose_type + "  : " + str(selected_choice) )
        c_type_count = type_map.get(choose_type)
        if c_type_count is None:
            type_map.add(choose_type, 1)
        else:
            type_map.add(choose_type, c_type_count + 1)

        type_file = open("destination/malware_grouped_result/" + str(choose_type), "a")
        type_file.write(hash_code + " \n")
        type_file.close()


type_map = HashMap()
for i in range(263):
    print("File------------------" + str(i))
    parse_vt_data("destination/malware_vt_result/" + str(i), type_map)
    print(type_map)